%%% this code plots the country graphs created by make_country_graphs.m

%clear
%close all
%clc

set(0,'DefaultFigureWindowStyle','normal','DefaultFigureVisible','on') % I stopped doing the docked thing because it doesn't allow to specify size and gives inconsistent results across computers

addpath('../Toolbox/export_fig/');
addpath('../Toolbox/struct2csv/');

graph_width=800;
graph_height=600;

% -----------------------
% DEFINE FOLDER LOCATIONS

folders;

% save?s
save_graphs=1;

% what to plot?
basic = 0;
plot_baseline = 0;
plot_cfactuals_var = 1;
plot_slides = 0;
grid = 0;
path_save_grid_plots = path_papers_graph;
osm = 0;
WP = 0;

%% load country list

country_list_short;

%% set discretization parameters to load the data
 % x_diag = 0;      % same, for diagonal paths
default_cellsize = 0.5;

% keep largest connected?
largest_connected = 1;
connected=largest_connected;

%% generate figure for Spain and France

for country_n = 1:Ncountries
% country_n=1;

    % Set default tolerance
    %% load map
    country_icc = char( countries(country_n) );  % country ICC code
    
    country = icc2name( country_icc );
    
    x_ver_hor = get_threshold( country ) ;   
    x_diag = x_ver_hor;
    x_thresh = x_ver_hor;
    
    disp( ['Country: ',country] )
    
    % load in the full grid mat 
    graph_filename = [ path_load_grids,country, '_grid_',...
                        num2str( x_diag ),'_',...
                        num2str( x_ver_hor ),'_',...
                        num2str( default_cellsize ),'_connected', ...
                        num2str( largest_connected ), ...
                        '_WP', num2str(WP), '_osm', num2str(osm),'.mat' ];
                    
    load( graph_filename );

    % unpack
    places_grid=country_graph.places_grid;
    gridmap=country_graph.gridmap;
    places2=country_graph.places2;
    discretized_roads=country_graph.discretized_roads;
    graph_export=country_graph.graph_export;
    unique_edges=country_graph.unique_edges;
    edges = country_graph.edges;
    unique_nodes=country_graph.unique_nodes;
    dist_info = country_graph.dist_info;
    
    % country bounds
    country_bounds=country_graph.country_bounds;
    
    %% compute some variables used in figures
    
    discretized_avI = cell2mat( {discretized_roads.avI} )';
    discretized_use = cell2mat( {discretized_roads.use} )';
    discretized_lanes = cell2mat( {discretized_roads.lanes} )';
    
    max_disc_avI = max( discretized_avI );
    
    keep_discretized = cell2mat( {discretized_roads.keep} )';  % these are the links that are kept
    
    %tabulate(discretized_avI)
    discretized_avI_int = round( discretized_avI );
    discretized_avI_int( ~keep_discretized ) = 0.1;
    %tabulate(discretized_avI_int)
    
    % common parameters to all graphs
    bright_line = 0.5;  % higher=brighter
    bright_nodes = 0.4; % higher=brighter
    max_bright = 5;    % higher=darker lines
    adj_width = 0.3;
    markersize = 5;
    color_growth = [ 0 1 0 ];   % red = [ 1 0 0 ]  green = [ 0 1 0 ] blue = [ 0 0 1 ]
    color_shrink = [ 1 0 0 ];
    
    % number of categories for plotting
    N_cat = numel( unique( cell2mat( {gridmap.quantiles} ) ) )-1;
    
    %%
    if plot_cfactuals_var
        
        adj_width = 6;
        max_bright = 1.5;
        GAMMA = 0.10;
%         CFAC_TYPE = {'exp_eng','mis_eng','exp_til'};
        CFAC_TYPE = {'exp_eng','mis_eng'}; %'',
        MOBILITY = {'off'}; % 'on',
        CONGESTION = {'on'};
        
        % parameters that remain constant throughout the loop across calibrations
        alpha = 0.4;
        beta = 0.13;
        sigma = 5;
        rho = 0;
        a = 1;
        Ngoods = 10;
        nu = 1;
        
        for cfac_type = CFAC_TYPE
            
            for cong = CONGESTION
                
                for mobil=MOBILITY
                    
                    for jj=1:length(GAMMA)
                        
                        gamma = GAMMA(jj);
                        
                        %% baseline parameters
                        param = init_parameters( 'a',a,'rho',rho,'alpha',alpha,'sigma',sigma,...       % preferences and technology
                            'beta',beta,'gamma',gamma,'nu',nu,'m',ones(Ngoods+1,1),...   % transpot costs
                            'K',1,...
                            'LaborMobility',char(mobil),...
                            'N',Ngoods+1,...
                            'CrossGoodCongestion',char(cong),...
                            'TolKappa',1e-4 );
                        
                        % define file to load calibration and set diary
                        filename = [ country,...
                            '_diag',num2str( x_diag ),...
                            '_hor',num2str( x_ver_hor ),...
                            '_a',num2str( param.a ),...
                            '_rho',num2str( param.rho ),...
                            '_alpha',num2str( param.alpha ),...
                            '_sigma',num2str( param.sigma ),...
                            '_beta',num2str( param.beta ),...
                            '_gamma',num2str( param.gamma ),...
                            '_nu',num2str( param.nu ),...
                            '_mobil',num2str( param.mobility ),...
                            '_cong',num2str( param.cong ),...
                            '_ngoods',num2str( Ngoods ),...
                            '_cellsize',num2str(default_cellsize), ...
                            '_WP', num2str(WP), ...
                            '_osm', num2str(osm)];
                        
                        % load
                        load( [  path_save_cfactuals ,filename,'_cfactual_', char(cfac_type),'.mat' ] )

                        % recover results and g
                        g = cfactual.g;
                        results_actual = cfactual.results_actual;
                        results_cfactual = cfactual.results_cfactual;
                        
                        % GDP per capita in traded sector in calibrated model, cfactual, and data
                        y_calib = sum( results_actual.Pjn.*results_actual.Yjn,2 )./results_actual.Lj;
                        y_exp = sum( results_cfactual.Pjn.*results_cfactual.Yjn,2 )./results_cfactual.Lj;
                        y_data = g.Y./g.L;
                        
                        % total GDP in calibration and cfactual
                        PHj_calib = (1-param.alpha) / param.alpha * results_actual.PCj .* results_actual.cj ./ cfactual.param.hj;
                        GDP_calib = sum(results_actual.Pjn.*results_actual.Yjn,2) + PHj_calib.*cfactual.param.Hj;
                        
                        PHj_exp = (1-param.alpha) / param.alpha * results_cfactual.PCj .* results_cfactual.cj ./ cfactual.param.hj;
                        GDP_exp = sum(results_cfactual.Pjn.*results_cfactual.Yjn,2) + PHj_exp.*cfactual.param.Hj;
                        
                        % population in calibrated model, cfactual, and data
                        L_calib = results_actual.Lj;
                        L_exp = results_cfactual.Lj;
                        L_data = g.L;
                        
                        % total income in calibrated model, cfactual, and data
                        Y_calib = L_calib.*y_calib;
                        Y_exp = L_exp.*y_exp;
                        Y_data = L_data.*y_data;
                        
                        % total consumption in calibrated and cfactual
                        C_calib = results_actual.Cj;
                        C_exp = results_cfactual.Cj;
                        
                        % per capita
                        c_calib = C_calib./L_calib;
                        c_exp = C_exp./L_exp;
                        
                        % welfare gain
                        welfare_gain = consumption_equivalent(cfactual.param,g,c_calib,L_calib,results_cfactual.welfare)-1;
                        
                        % exports and imports in calib
                        [ X_calib,M_calib ] = recover_X_M( results_actual,g );
                        
                        % exports and imports in cfactual
                        [ X_exp,M_exp ] = recover_X_M( results_actual,g );
                        
                        % X/GDP
                        XY_calib = X_calib./Y_calib;
                        XY_exp = X_exp./Y_exp;
                        
                        % M/GDP
                        MY_calib = M_calib./Y_calib;
                        MY_exp = M_exp./Y_exp;
                        
                        % net exports over GDP
                        nx_calib = ( X_calib-M_calib )./Y_calib;
                        nx_exp = ( X_exp-M_exp )./Y_exp;
                        
                        % fundamentals
                        H_calib = cfactual.param.Hj;
                        Z_calib = max( cfactual.param.Zjn,[],2 );
                        
                        % infrastructure long
                        avI_obs = triu(g.avI);
                        avI_obs = avI_obs(:);
                        
                        avI_exp = triu(results_cfactual.Ijk);
                        avI_exp = avI_exp(:);
                        
                        % quantiles in calibrated and counterfactual model
                        quantile_L_calib = assign_quantiles( L_calib );
                        quantile_L_exp = assign_quantiles( L_exp );
                        
                        % lanes in calibrated and counterfactual model
                        I_obs = zeros( length(unique_edges),1 );
                        I_exp = zeros( length(unique_edges),1 );
                        for j=1:length(unique_edges)
                            
                            I_obs(j) = g.avI( unique_edges(j,1),unique_edges(j,2) );
                            I_exp(j) = results_cfactual.Ijk( unique_edges(j,1),unique_edges(j,2) );

                            
                        end
                        
                        % identify producers of differentiated goods
                        diff_producers = find( results_actual.Yjn(:,1)==0 );  % indexes of differentiated producers
                        
                        % compute expansion
                        gL = ( L_exp-L_calib )./( 0.5*( L_exp+L_calib ) );
                        gc = ( c_exp-c_calib )./( 0.5*( c_exp+c_calib ) );
                        gy = ( y_exp-y_calib )./( 0.5*( y_exp+y_calib ) );
                        gXY = ( XY_exp-XY_calib )./( 0.5*( XY_exp+XY_calib ) );
                        gMY = ( MY_exp-MY_calib )./( 0.5*( MY_exp+MY_calib ) );
                        gnx = ( nx_exp-nx_calib )./( 0.5*( nx_exp+nx_calib ) );
                        dI = I_exp-I_obs;
                        dL =  L_exp-L_calib;
                        gL2 = ( L_exp-L_calib )./( L_calib  );
                        gI = (I_exp-I_obs)./((I_exp+I_obs)*0.5);
                        
                        if strcmp(cfac_type, 'mis_eng') && ~strcmp(country, 'Mexico')
                            dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                            dI(dI<prctile(dI, 5)) = prctile(dI, 5);
                        
                        elseif strcmp(cfac_type, 'mis_eng') && strcmp(country, 'Mexico')
                            dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                            dI(dI<prctile(dI, 5)) = prctile(dI, 5);
                        end
                        
                        if strcmp(cfac_type, 'exp_eng')
                            dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                        end

                        % max I
                        maxI_obs = max( max(I_obs),max(I_exp) );
                        %[ I_exp( I_obs==0.0001 ) I_obs( I_obs==0.0001 ) I_exp( I_obs==0.0001 )./I_obs( I_obs==0.0001 ) ]
                        
                        show_cells = 0;  % otherwise, show nodes
                        
                        %% plot change in network and in population
                        
                        close(figure(6))
                        h=figure(6);
                        set(h,'Position',[0 0 graph_width graph_height]);

                        % For countries with grids that we made shorter, 
                        % plot the entire boundary 
                        if strcmp(country, 'Peru') || strcmp(country, 'Chile')

                            file = [datafolder_ne,'ne_10m_admin_0_map_units.shp'];
                            country_bounds = shaperead(file,'Selector',{@(name)strcmp( name,country ),'GEOUNIT'}); 


                        end


                        mapshow( country_bounds,...
                            'FaceColor','Black',...
                            'EdgeColor','White','LineWidth', 0.5 )


                            % population
                            x = dL;
                            x(x<prctile(x,3.5)) = prctile(x,3.5);
                            x(x>prctile(x,96.5)) = prctile(x,96.5);
                            if show_cells && param.mobility
                                for i=1:length( gridmap )

                                    norm_neg = 1/abs( mean(x(x<0)) )*1/2.5;   %gL(gL<0)*norm_neg
                                    norm_pos = 1/max( x(x>0) )*2.5;         %gL(gL>0)*norm_neg

                                    relsize = x(i);
                                    bring_fire = 0;

                                    color = min( max( -relsize*norm_neg,0 )*color_shrink+...
                                        ( max( -relsize*norm_neg,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 )+...
                                        min( max( relsize*norm_pos,0 )*color_growth+...
                                        ( max( relsize*norm_pos,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 );

                                    mapshow( gridmap( i ),...
                                        'FaceColor',color );

                                end
                            end

                            % lanes
                            for j=1:length( dI )
                                [ j ]
                                % COLOR SCALE OF LANES
    %                             colorweight = abs( dI(j) )/max_bright;
                                colorweight = abs( dI(j) )/max(abs(dI))/max_bright;
                                lanewidth = abs( dI(j) )/max(abs(dI))*adj_width;

                                if dI(j)>0
                                    color =  min( bright_line+( 1-bright_line )*colorweight*color_growth,1 );
                                elseif dI(j)<0
                                    color =  min( bright_line+( 1-bright_line )*colorweight*color_shrink,1 );
                                end

                                %if discretized_roads( j ).keep==1
                                exclude = [15, 354, 346, 362, 350];
                                if (ismember(j, exclude) == 0 || ~strcmp(country, 'Mexico')) && dI(j) ~= 0
                                        mapshow( discretized_roads( j ),...
                                        'Color',min( color,1 ),...
                                        'LineWidth',lanewidth )
                                end

                                %end
                                hold on;

                            end
                            
                            
                            if ~show_cells
                                for i=1:length( gridmap )
                                    [i ]
                                    norm_neg = 1/abs( mean(x(x<0)) )*1/2.5;   %gL(gL<0)*norm_neg
                                    norm_pos = 1/max( x(x>0) )*2.5;         %gL(gL>0)*norm_neg

                                    relsize = x(i);
                                    bring_fire = 0;

                                    if param.mobility
                                        color = min( max( -relsize*norm_neg,0 )*color_shrink+...
                                            ( max( -relsize*norm_neg,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 )+...
                                            min( max( relsize*norm_pos,0 )*color_growth+...
                                            ( max( relsize*norm_pos,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 );
                                    else
                                        color = [ 0 0 0 ];
                                    end
                                    
                                    mapshow( places2( i ),...
                                    'Marker','o',...
                                    'MarkerFaceColor',color,...
                                    'MarkerEdgeColor',[ 1 1 1 ],...
                                    'MarkerSize',markersize )

                                end
                            end

                            margin = 0.5;
                            Xlim = [ min( cell2mat({country_bounds.X}) )-margin;max( cell2mat({country_bounds.X}) )+margin ];
                            Ylim = [ min( cell2mat({country_bounds.Y}) )-margin;max( cell2mat({country_bounds.Y}) )+margin ];
                            set( gca,'XTick',[],'YTick',[],'XLim',Xlim,'YLim',Ylim,'Box','on' )

                            if save_graphs
                               export_fig([ path_papers_graph,filename,'_cfactual_dL_dI_',char(cfac_type),'.eps' ], '-depsc');
                            end
                    end
                                                  
                    end

                end

            end

        end

    end


